package com.hj.aiagent.core.flow;

import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.hj.aiagent.core.agent.BaseAgent;
import com.hj.aiagent.core.agent.ToolCallAgent;
import com.hj.aiagent.core.agent.model.AgentState;
import com.hj.aiagent.core.flow.model.Plan;

import com.hj.aiagent.core.flow.model.PlanStatus;
import com.hj.aiagent.core.flow.model.PlanStep;
import com.hj.aiagent.core.flow.model.StepStatus;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.extern.slf4j.Slf4j;

import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.tool.ToolCallback;

import java.util.ArrayList;

import java.util.List;
import java.util.UUID;

/**
 * 规划并执行模式的智能体
 */
@EqualsAndHashCode(callSuper = true)
@Data
@Slf4j
public class PlanningAgent extends BaseAgent {
    
    // 当前执行的计划
    private Plan currentPlan;
    
    // 可用工具
    private final ToolCallback[] availableTools;
    
    // 提示词模板
    private static final String PLANNING_PROMPT = """
        你是一个任务规划专家。请将用户的复杂任务分解为一个结构化的执行计划。
        
        你的输出必须是一个有效的JSON对象，包含以下字段：
        {
            "summary": "计划的简要总结",
            "steps": [
                {
                    "id": "step_1",
                    "goal": "这个步骤要实现的目标",
                    "description": "详细的步骤描述",
                    "expectedTool": "预期使用的工具名称(可选)",
                    "dependencies": ["step_0"] // 依赖的步骤ID列表(可选)
                },
                ...
            ]
        }
        
        示例任务："规划5天云南旅游"
        示例输出：
        {
            "summary": "为用户规划一个5天的云南旅游行程，包括景点筛选、行程安排和PDF文档生成",
            "steps": [
                {
                    "id": "step_1",
                    "goal": "筛选云南热门景点",
                    "description": "根据5天时间限制，筛选出云南最值得游览的景点",
                    "expectedTool": "filterAttractions"
                },
                {
                    "id": "step_2",
                    "goal": "制定详细行程",
                    "description": "根据筛选的景点，安排合理的游览顺序和时间",
                    "expectedTool": "createItinerary",
                    "dependencies": ["step_1"]
                },
                {
                    "id": "step_3",
                    "goal": "生成PDF文档",
                    "description": "将行程安排生成PDF文档供用户下载",
                    "expectedTool": "generatePDF",
                    "dependencies": ["step_2"]
                }
            ]
        }
        """;
    
    private static final String EXECUTION_PROMPT = """
        你正在执行计划的当前步骤。请根据步骤目标完成相应任务。
        
        当前步骤信息：
        目标: %s
        描述: %s
        
        请完成这个步骤的任务。
        """;
    
    private static final String EVALUATION_PROMPT = """
        请评估当前步骤的执行结果，并决定下一步行动。
        
        原始目标: %s
        当前步骤: %s
        执行结果: %s
        
        请输出一个JSON对象，包含以下字段：
        {
            "decision": "CONTINUE|REPLAN|FINISH",  // 决策：继续下一步、重新规划或完成任务
            "reason": "决策的原因",
            "nextStepThoughts": "对下一步的思考(可选)"
        }
        """;
    
    public PlanningAgent(ToolCallback[] availableTools) {
        super();
        this.availableTools = availableTools;
    }
    
    @Override
    public String step() {
        // 1. 如果没有计划，先创建计划
        if (currentPlan == null) {
            return createPlan();
        }
        
        // 2. 如果计划已完成，返回结果
        if (currentPlan.isCompleted()) {
            setState(AgentState.FINISHED);
            return "任务完成！\n" + generatePlanSummary();
        }
        
        // 3. 执行当前步骤
        PlanStep currentStep = currentPlan.getCurrentStep();
        if (currentStep == null) {
            setState(AgentState.FINISHED);
            return "所有步骤已完成";
        }
        
        // 4. 执行步骤
        String stepResult = executeStep(currentStep);
        currentStep.setResult(stepResult);
        currentStep.setStatus(StepStatus.COMPLETED);
        
        // 5. 评估执行结果
        String evaluation = evaluateExecution(currentStep);
        
        // 6. 根据评估结果决定下一步
        if (evaluation.contains("REPLAN")) {
            currentPlan.setStatus(PlanStatus.NEED_REPLAN);
            return "需要重新规划: " + evaluation;
        } else if (evaluation.contains("FINISH")) {
            setState(AgentState.FINISHED);
            currentPlan.setStatus(PlanStatus.COMPLETED);
            return "任务提前完成: " + evaluation;
        } else {
            // 继续下一步
            currentPlan.moveToNextStep();
            return "步骤完成: " + currentStep.getGoal() + "\n" + evaluation;
        }
    }
    
    /**
     * 创建执行计划
     */
    private String createPlan() {
        try {
            // 获取用户原始请求
            String userGoal = getMemory().get(0).getText();
            
            // 调用LLM生成计划
            Prompt prompt = new Prompt(PLANNING_PROMPT);
            ChatResponse response = getClient().prompt(prompt)
                .system("你是一个任务规划专家")
                .user("请为以下任务创建执行计划: " + userGoal)
                .call()
                .chatResponse();
            
            // 解析LLM返回的计划
            String planJson = response.getResult().getOutput().getText();
            JSONObject planObj = JSONUtil.parseObj(planJson);
            
            // 创建计划对象
            String planId = UUID.randomUUID().toString();
            Plan plan = new Plan(planId, userGoal);
            plan.setSummary(planObj.getStr("summary"));
            
            // 解析步骤
            List<PlanStep> steps = new ArrayList<>();
            for (Object stepObj : planObj.getJSONArray("steps")) {
                JSONObject stepJson = (JSONObject) stepObj;
                PlanStep step = new PlanStep(
                    stepJson.getStr("id"),
                    stepJson.getStr("goal"),
                    stepJson.getStr("description")
                );
                step.setExpectedTool(stepJson.getStr("expectedTool"));
                if (stepJson.containsKey("dependencies")) {
                    step.setDependencies(stepJson.getBeanList("dependencies", String.class));
                }
                steps.add(step);
            }
            plan.setSteps(steps);
            
            // 设置当前计划
            this.currentPlan = plan;
            currentPlan.setStatus(PlanStatus.IN_PROGRESS);
            
            return "已创建执行计划:\n" + formatPlan(plan);
            
        } catch (Exception e) {
            log.error("创建计划失败", e);
            setState(AgentState.ERROR);
            return "创建计划失败: " + e.getMessage();
        }
    }
    
    /**
     * 执行单个步骤
     */
    private String executeStep(PlanStep step) {
        try {
            step.setStatus(StepStatus.IN_PROGRESS);
            
            // 如果指定了工具，使用工具调用代理执行
            if (StrUtil.isNotBlank(step.getExpectedTool())) {
                ToolCallAgent toolAgent = new ToolCallAgent(availableTools);
                toolAgent.setClient(getClient());
                toolAgent.setMemory(getMemory());
                
                String stepPrompt = String.format(EXECUTION_PROMPT, step.getGoal(), step.getDescription());
                return toolAgent.run(stepPrompt);
            } else {
                // 否则直接调用LLM
                Prompt prompt = new Prompt(String.format(EXECUTION_PROMPT, step.getGoal(), step.getDescription()));
                ChatResponse response = getClient().prompt(prompt).call().chatResponse();
                return response.getResult().getOutput().getText();
            }
        } catch (Exception e) {
            log.error("执行步骤失败: " + step.getGoal(), e);
            step.setStatus(StepStatus.FAILED);
            return "执行失败: " + e.getMessage();
        }
    }
    
    /**
     * 评估执行结果
     */
    private String evaluateExecution(PlanStep step) {
        try {
            Prompt prompt = new Prompt(String.format(
                EVALUATION_PROMPT,
                currentPlan.getOriginalGoal(),
                step.getGoal(),
                step.getResult()
            ));
            
            ChatResponse response = getClient().prompt(prompt).call().chatResponse();
            String evaluationJson = response.getResult().getOutput().getText();
            JSONObject evalObj = JSONUtil.parseObj(evaluationJson);
            
            String decision = evalObj.getStr("decision");
            String reason = evalObj.getStr("reason");
            
            return String.format("决策: %s, 原因: %s", decision, reason);
        } catch (Exception e) {
            log.error("评估执行结果失败", e);
            return "评估失败，继续下一步";
        }
    }
    
    /**
     * 格式化计划输出
     */
    private String formatPlan(Plan plan) {
        StringBuilder sb = new StringBuilder();
        sb.append("计划ID: ").append(plan.getId()).append("\n");
        sb.append("目标: ").append(plan.getOriginalGoal()).append("\n");
        sb.append("总结: ").append(plan.getSummary()).append("\n");
        sb.append("步骤:\n");
        
        for (int i = 0; i < plan.getSteps().size(); i++) {
            PlanStep step = plan.getSteps().get(i);
            sb.append(String.format("  %d. %s [%s]\n", 
                i+1, step.getGoal(), step.getStatus()));
            sb.append(String.format("     描述: %s\n", step.getDescription()));
            if (StrUtil.isNotBlank(step.getExpectedTool())) {
                sb.append(String.format("     工具: %s\n", step.getExpectedTool()));
            }
        }
        
        return sb.toString();
    }
    
    /**
     * 生成计划总结
     */
    private String generatePlanSummary() {
        StringBuilder sb = new StringBuilder();
        sb.append("计划执行完成!\n");
        sb.append("原始目标: ").append(currentPlan.getOriginalGoal()).append("\n");
        sb.append("执行结果:\n");
        
        for (PlanStep step : currentPlan.getSteps()) {
            sb.append(String.format("- %s: %s\n", 
                step.getGoal(), 
                step.getStatus() == StepStatus.COMPLETED ? "成功" : "失败"));
            if (StrUtil.isNotBlank(step.getResult())) {
                sb.append("  结果: ").append(step.getResult()).append("\n");
            }
        }
        
        return sb.toString();
    }
}
